#include "StdAfx.h"
#include "Limitor.h"
#include "Tlhelp32.h"



CLimitor::CLimitor(void)
{
	watch=true;
    targetPid=-1;
	realTime=false;
	watchThread=nullptr;
	hackThread=nullptr;
	timeGreen=67;
	watchSemaphore= gcnew Semaphore(1,1);
}

void CLimitor::Start()
{
	watch=true;
	watchThread= gcnew Thread(gcnew ThreadStart(this,&CLimitor::Watch));
	watchThread->Start();
}

void CLimitor::Stop(void)
{
	watch=false;
	//if(hackThread != nullptr)
	//{
	//	hackThread->Join();
	//}
	if(watchThread != nullptr)
	{
		watchThread->Join();
	}
}

void CLimitor::SetTarget(int pid)
{
	if(targetPid!=-1)
	{
		Stop();
	}
	targetPid=pid;
}

void CLimitor::Watch()
{
	while(watch)
	{
		PostMessage(m_HWMD,WM_USER_STOP,(WPARAM)targetPid,(LPARAM)START);
		watchSemaphore->WaitOne();
		hackThread=gcnew Thread(gcnew ThreadStart(this,&CLimitor::Hack));
		hackThread->Start();
	}
	PostMessage(m_HWMD,WM_USER_STOP,(WPARAM)targetPid,(LPARAM)STOP);
}

void CLimitor::Hack()
{
	BOOL bToken = FALSE;
	BOOL bPriv = FALSE;
	HANDLE hToken = NULL;

	bToken = OpenThreadToken( GetCurrentThread(), 
					TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, 
					TRUE, &hToken );
	if( bToken )
	{
		WriteDebug ("OpenThreadToken OK" );
	}
	else
	{
		if( ERROR_NO_TOKEN == GetLastError() )
		{
			WriteDebug( "ERROR_NO_TOKEN"  );
		}
		else
		{
			WriteDebug ( "OpenThreadToken failed, Error Code" );
		}
		BOOL bImp = ImpersonateSelf(SecurityImpersonation);
		if( bImp )
		{
			bToken = OpenThreadToken(GetCurrentThread(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, FALSE, &hToken);
		}
		else
		{
			WriteDebug ("ImpersonateSelf failed" );
		}
	}

	if( bToken )
	{
		WriteDebug( "OpenThreadToken OK" );
		bPriv = EnableDebugPrivilege( hToken, TRUE );
		if( bPriv ){
			WriteDebug("EnableDebugPrivilege Success!");
		}
		else {
			WriteDebug("EnableDebugPrivilege Failed!");
		}
	}

	HANDLE hProcess = OpenProcess(
				PROCESS_QUERY_INFORMATION | PROCESS_SET_INFORMATION ,
				FALSE,
				targetPid
			);
	if( hProcess )
	{
		DWORD dwOldPriority = GetPriorityClass( hProcess );
		if(realTime )
		{
			 SetPriorityClass( hProcess, IDLE_PRIORITY_CLASS );
		}
	}
	else
	{
		WriteDebug("Open process failed!");
		CloseHandle( hToken );
	    watchSemaphore->Release(1);
		watch=false;
        PostMessage(m_HWMD,WM_USER_STOP,(WPARAM)targetPid,(LPARAM)TARGET_MISSING);
		return;
	}

	DWORD dwThreadId[ MAX_THREADS ];
	int iThreadCount = ListProcessThreads( this->targetPid, dwThreadId );
	
	if( iThreadCount == 0 )
	{
		CloseHandle( hToken );
		CloseHandle( hProcess );
		watchSemaphore->Release(1);
		return;
	}

	HANDLE hTarget[ MAX_THREADS ];
	ZeroMemory( hTarget, sizeof( HANDLE ) * MAX_THREADS );
	int iOpenedThread = 0;
	for( int i = 0; i < iThreadCount; i++ )
	{
		hTarget[ i ] = OpenThread( THREAD_SUSPEND_RESUME, FALSE, dwThreadId[ i ] );
		if( hTarget[ i ] == NULL )
		{
			WriteDebug("Open thread failed.");
		}
		else
		{
			iOpenedThread++;
		}
	}

	
	for( DWORD dwInnerLoops = 1; dwInnerLoops <= MAX_LOOPS; dwInnerLoops++ )
	{
		BYTE TimeRed = (BYTE)( 100 - timeGreen );
		BYTE TimeGreen = timeGreen;
		
		for( int i = 0; i < iThreadCount; i++ )
		{
			if( ! hTarget[ i ] ) continue;

			DWORD dwPrevSuspendCount = SuspendThread( hTarget[ i ] );
			if( ( DWORD ) -1 == dwPrevSuspendCount )
			{
				WriteDebug("Suspend thread failed");
			}
		}

		Sleep( (DWORD) TimeRed );
		for( int i = 0; i < iThreadCount; i++ )
		{
			if( ! hTarget[ i ] ) continue;
			DWORD dwPrevSuspendCount = ResumeThread( hTarget[ i ] );
			if( ( DWORD ) -1 == dwPrevSuspendCount )
			{
				WriteDebug(" ResumeThread failed");
			}
		}

		Sleep( (DWORD) TimeGreen );
		if(!watch) break;
	}

	for( int i = 0; i < iThreadCount; i++ )
	{
		if( hTarget[ i ] ) CloseHandle( hTarget[ i ] );
	}
	watchSemaphore->Release(1);
}

BOOL CLimitor::EnableDebugPrivilege( HANDLE hToken, BOOL bEnable )
{
	TOKEN_PRIVILEGES tp;
    LUID luid;

    if( ! LookupPrivilegeValue( NULL, SE_DEBUG_NAME, &luid ) ) return FALSE;

    tp.PrivilegeCount = 1UL;
    tp.Privileges[ 0 ].Luid = luid;
    tp.Privileges[ 0 ].Attributes = 0UL;

    TOKEN_PRIVILEGES tpPrev;
    DWORD cbPrev = sizeof( TOKEN_PRIVILEGES );
    
	AdjustTokenPrivileges( hToken, FALSE, &tp, sizeof(TOKEN_PRIVILEGES),
            &tpPrev, &cbPrev );
    
	if ( GetLastError() != ERROR_SUCCESS )
	{
		return FALSE;
	}

    tpPrev.PrivilegeCount = 1;
    tpPrev.Privileges[ 0 ].Luid = luid;

    if( bEnable )
	{
        tpPrev.Privileges[ 0 ].Attributes |= SE_PRIVILEGE_ENABLED;
    }
    else
	{
        tpPrev.Privileges[ 0 ].Attributes &= ~SE_PRIVILEGE_ENABLED;
    }

    AdjustTokenPrivileges( hToken, FALSE, &tpPrev, cbPrev, NULL, NULL );

    if ( GetLastError() != ERROR_SUCCESS )
	{
		return FALSE;
	}
    return TRUE;
}

int CLimitor::ListProcessThreads( DWORD dwOwnerPID, DWORD * dwThreadIdTable ) 
{
	int iThreadCount = 0; 
	HANDLE hThreadSnap = CreateToolhelp32Snapshot( TH32CS_SNAPTHREAD, 0UL );
	if( hThreadSnap == ( HANDLE ) -1 ) return 0;

	THREADENTRY32 te32;
	te32.dwSize = sizeof( THREADENTRY32 ); 

	if( ! Thread32First( hThreadSnap, &te32 ) )
	{
		CloseHandle( hThreadSnap );
		return 0;
	}

	do 
	{ 
		if( te32.th32OwnerProcessID == dwOwnerPID )
		{
			dwThreadIdTable[ iThreadCount++ ] = te32.th32ThreadID;
		}
	} while( Thread32Next( hThreadSnap, &te32 ) && iThreadCount < MAX_THREADS ); 

	CloseHandle( hThreadSnap );
	return iThreadCount;
}

void CLimitor::WriteDebug(String^ message)
{
	Debug::WriteLine(message);
	Console::WriteLine(message);
}

void CLimitor::SetTimeGreen(int value)
{
	timeGreen=value;
}
